import scipy
from scipy import asarray, randn, r_, pi, arange, sqrt
from scipy.special import erf
from scipy.optimize import leastsq
from scipy import poly1d

def erf_fit(y, x=None, guess_p = [0., 1., 1., 0.]):
    def residuals(p, y, x):
        dcGain, height, tau, shift = p
        err = y - (dcGain + height * erf( (x - shift) / float(tau) ))
        return err

    def func_eval(e, p):
        print 'p ',p
        return p[0] + p[1] * erf( (e - p[3]) / float(p[2]))
    
    if x == None:
        x = asarray(r_[0:len(y)], typecode = 'd')
    foo = list(leastsq( residuals, guess_p, args=(y, x), maxfev=10**6))
    foo.append( lambda x: func_eval(x, foo[0]))
    foo.append( lambda x: foo[0][0] + \
                (x - foo[0][3]) * 2 * foo[0][1] / (sqrt(pi) * foo[0][2]))
    foo.append( 2 * foo[0][1] / (sqrt(pi) * foo[0][2]))
    return foo

def maxErfDerivative(y, x=None):
    """Takes a set of points, y, and returns the maximum derivate via
    an erf funciton."""
    foo = erf_fit(y, x)
    return foo[0][1] / foo[0][2] * 2. / sqrt(pi)

def polyFit(y, x=None, degree=2, guess_coef = None, maxIter = 1000):
    """Takes an array y, and an optional x array, and returns a polynomial
    of degree N using a least squares fit - i.e. a polynomial close to:

    y[i] = a_0 * x[i] ** N + a_1 * x[i] ** (N-1) + ... + a_N

    inputs:
        y - target function
        x - index vector - if None, uses arange(len(y))
        degree - degree of polynomial to return
        guess_coef - guess for coefficients, len(guess_coef) == len(y)
        maxIter - maximum iterations for least squares function

    outputs a list containing (in order):
       an instance of scipy_base.poly1d
       information returned from leastsq
    """
    def residuals(p, y, x):
        poly = poly1d(p)
        return y - poly(x)

    if x==None:
        x = arange(len(y), typecode = 'd')
    if guess_coef == None:
        guess_coef = asarray([1.] * (degree + 1))
    if len(x) != len(y):
        raise 'LengthError','Index Vector and Target Vector'+\
              'need to have the same length.'
    if degree + 1 != len(guess_coef):
        #print 'degree',degree
        #print 'guess coefs',guess_coef
        raise 'LengthError','Must have (degree + 1) guess coefs.'

    tmp = leastsq( residuals, guess_coef, args=(y, x), maxfev = maxIter )

    return [poly1d(tmp[0]), tmp[1]]

def test(pl=False):
    t = r_[-10:100:.1]
    y_true = 1.15 + - .04 * erf( (t - 30.) / 30.333 )
    y_meas = y_true + .001*randn(len(y_true))
    #y_meas = y_true + map(lambda x: (abs(x + 4) < 1 and .01) or 0, t)

    fit_values = erf_fit(y_meas, t, [3.5, 4.5, 3.0, 0])

    if pl:
        from pylab import plot, show
        plot(t, y_meas, 'bo')
        plot(t, y_true, 'g+')
        plot(t, fit_values[2](t), 'r-')
        show()

    return erf_fit(y_meas)#, t, [3.5, 4.5, 3.0])
